Repository navigation
perf: don't import the Python HTTP and OAuth client stack in connect() - #1051
vuanhphung wants to merge 7 commits into
Conversation
common/http.py imported requests, urllib3 and the retry policy without using them, auth/common.py imported the retry policy without using it, and auth/authenticators.py imported the OAuth client (oauthlib, jwt) at module load for two code paths that need it. Every connect() paid for these in its first import. Drop the unused imports and import the OAuth classes where they are used. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
GitHub Actions
Error
GitHub Actions
Error
telemetry_client imported the telemetry HTTP client (urllib3) and push clients (pybreaker) at module load, and telemetry_client and the kernel auth bridge imported token federation (jwt) for isinstance checks; client.py, session.py and the cloudfetch downloader imported UnifiedHttpClient at module load. Import the HTTP and push clients where they are constructed, check for token federation through sys.modules (an instance can't exist before its module is imported), and resolve client.UnifiedHttpClient lazily so its patch seam keeps working. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
GitHub Actions
Inactive
GitHub Actions
Inactive
vuanhphung
changed the title
GitHub Actions
Error
![peco-review-bot[bot]](https://avatars.githubusercontent.com/in/3801177?s=60&v=4){kind=small}
There was a problem hiding this comment.
Verdict: 1 Low
Looks good, with one low concern: the new import-guard test can skip when it should fail. I checked the removed module-level names (common.http, telemetry_client, authenticators, session, downloader). Nothing in src/ or tests/ imports or patches them, and the client.__getattr__ route still supports patch("databricks.sql.client.UnifiedHttpClient").
Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
GitHub Actions
Error
GitHub Actions
Inactive
GitHub Actions
Error
GitHub Actions
Error
There was a problem hiding this comment.
Verdict: 1 Low
Looks good: one low-severity concern. Some names that were importable at module level before are now silently gone, outside the client.py __getattr__ shim. The deferred imports, the sys.modules-based isinstance check and the subprocess tests with sentinel exit codes look correct. Nit: in authenticators.py the if TYPE_CHECKING: block sits in the middle of the runtime import list; moving it after the last import would read better.
GitHub Actions
Error
The kernel backend re-raises a missing databricks_sql_kernel as an ImportError, so look for the missing extra along the exception chain. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
GitHub Actions
Error
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
authenticators, auth, telemetry_client, session and the cloudfetch downloader stopped importing OAuthManager, ClientCredentialsTokenSource, RefreshableTokenSource, TokenFederationProvider, UnifiedHttpClient and the telemetry push clients at module load. Downstream imports or patch()es of those module attributes would then fail. Each module now resolves them on first access (PEP 562), as client.py already does, and the places that use them read the module attribute, so a patch there still takes effect. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
GitHub Actions
Error
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Error
There was a problem hiding this comment.
Verdict: 1 Nit
Looks good. I found no correctness issues. The dropped imports in common/http.py, auth/common.py and cloudfetch/downloader.py have no remaining users in the repo. The PEP 562 __getattr__ re-exports cover every moved name, so imports and mock.patch keep working. The sys.modules-based token-federation check is sound, because an instance can't exist until its module has been imported.
Other findings
- ⚪ Nit — Nit: the lazy-name pattern is written three ways:
importlib.import_modulewith a_LAZY_NAMESmap intelemetry_client.py, anif name in (...)tuple inauth.py,authenticators.pyanddownloader.py, andif name == ...inclient.pyandsession.py. One shared helper would make it easier to keep the names in sync. Also,dir()on these modules no longer lists the moved names, which matters only for introspection.
auth, authenticators and the cloudfetch downloader now use the same _LAZY_NAMES map as telemetry_client instead of if-chains. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
GitHub Actions
Inactive
GitHub Actions
Inactive
|
On the lazy-name Nit: in cbdae71, |
GitHub Actions
Error
There was a problem hiding this comment.
Verdict: 1 Low
Looks good. The lazy-import refactor is correct: the dropped imports in common/http.py and auth/common.py had no remaining users, the PEP 562 __getattr__ keeps every previously exposed name importable and patchable, and the sys.modules check for token federation is sound. One low-severity test-coverage gap. Nit: in authenticators.py, import importlib comes before import abc, and the TYPE_CHECKING block sits in the middle of the import list.
The new tests reached OAuthManager and the plain telemetry push client, but not get_token_source's ClientCredentialsTokenSource or the circuit-breaker branch of TelemetryClient, so a mistyped lazy-name key there would only fail at runtime. Co-authored-by: Isaac <no-reply@databricks.com> Signed-off-by: Vu Anh Phung <vu.phung@databricks.com>
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Error
GitHub Actions
Inactive
What type of PR is this?
Description
The first
connect()in a process importsdatabricks.sql.client, and with it the whole Python HTTP and OAuth client stack:requests,urllib3,oauthlib,jwtandpybreaker. That happens even when nothing uses them. A kernel connection (use_kernel=True) never does, because the kernel owns HTTP, auth and telemetry export. Most of these imports were dead or only needed further down one code path:common/http.pyimportedrequests,HTTPAdapter, urllib3'sRetryandDatabricksRetryPolicywithout using them, andauth/common.pyimportedDatabricksRetryPolicywithout using it.cloudfetch/downloader.pyimportedUnifiedHttpClientwithout using it. These are dropped.auth/authenticators.pybuiltOAuthManager/ClientCredentialsTokenSource(oauthlib, jwt) in two providers, andauth/auth.pywraps providers inTokenFederationProvider(jwt). These are now imported where they are constructed.telemetry_client.pyimportedUnifiedHttpClient(urllib3) and the push clients (pybreaker), which onlyTelemetryClient.__init__uses, so they are imported there. It and the kernel auth bridge also importedTokenFederationProviderforisinstancechecks. Those checks now go throughsys.modules: an instance can't exist before its module is imported.client.pyresolvesUnifiedHttpClientlazily through the module__getattr__it already uses for the Thrift names, sopatch("databricks.sql.client.UnifiedHttpClient")keeps working.session.pyonly needed it for an annotation.Behavior is unchanged: Thrift and SEA connections import what they use when they construct it. After this, a kernel connect and query import none of these modules. Related: #1049 defers pandas the same way.
How is this tested?
New tests import
databricks.sql,client,session,auth.auth,auth.authenticators,common.http, the cloudfetch downloader,telemetry_clientand the kernel client in fresh interpreters, and assert that none of the five packages loads. They fail onmain. More tests check that the token-federation check doesn't import token federation, still recognizes a realTokenFederationProvider(telemetry still reports a wrapped PAT as PAT), and that the browser OAuth provider still builds itsOAuthManager. The full unit suite, therealkernelset and mypy pass.Fresh processes, interleaved, median of 5 (import) or 3 (burst):
main(0c1b9e8)import databricks.sql.client, site-packages with.pyc.pyc(simulated)import databricks.sql.clientconnect(use_kernel=True), local mock with 1 ms latency:connect()p50In a separate isolated live run (4 processes × 200 sessions from a cold start, 16 cores,
mainvs this PR), each session's open time measured from process launch had a p50 of 1.68 s instead of 1.91 s, peak RSS was 9–11 MB lower per process, and QPS was unchanged."Without
.pyc" meansPYTHONDONTWRITEBYTECODE=1with a bytecode cache prefix that holds only the stdlib. That is what an install without compiled bytecode looks like, e.g.uv pip installwithout--compile-bytecodeinto a location the process can't write.Live benchmark (staging Large Reyden warehouse)
TPC-H Q6 on a Large Real-Time warehouse, 800 sessions (4 processes x 200 threads), cold start, each session querying back to back for 120 s, 16 client cores. Isolated A/B: only this change vs connector main 0c1b9e8 + kernel 1.4.0, 3 interleaved rounds, nothing else on the warehouse. Measured at 8320850. Median (min-max):
QPS within noise; small connect-time and memory gain.
Related Tickets & Documents
RE-17633 (connection-creation time at high concurrency).
This pull request and its description were written by Isaac.